#pragma once

#include "BST.h"


template <typename T, typename NodeType=BTNode<T>>
class AVLT: public BST<T, NodeType>
{
public:
	using Node = NodeType;
	/**
	 * 插入元素
	 * @param element 
	 */
	void Insert(const T& element);

	/**
	 * 删除节点
	 * @param deletedNode 
	 */
	void Delete(Node* deletedNode);
private:
	// 检查二叉树平衡，如果平衡被破坏，修复平衡
	void OnInsert_MaintainBalance(const T& element);

	void OnDeleted_MaintainBalance(Node* childRoot);

	// 平衡二叉树
	void DoBSTBalance(Node* parent, Node* unbalancedSubtree, const T& element, const int balanceFactor);
	
	// 搜索最小不平衡左子树
	Node* FindUnbalancedSubtree(Node* childRoot, const T& element);

	/**
	 * 左旋
	 * @param childRoot 最小不平衡子树根节点
	 * @return 返回旋转后的新子树的根节点
	 */
	Node* LeftRotate(Node* childRoot);

	/**
	 * 右旋
	 * @param childRoot 最小不平衡子树根节点
	 * @return 返回旋转后的新子树的根节点
	 */
	Node* RightRotate(Node* childRoot);

	// 找到不平衡子树根节点的最深叶节点
	Node* FindDeepestLeaf(Node* childRoot);
	Node* FindDeepestLeafRecursively(Node* origRoot, Node* childRoot, const int rootHeight);

	// 计算到指定节点的深度
	int CalculateDepth(Node* targetNode, Node* calculatedNode);
};

template <typename T, typename NodeType>
void AVLT<T, NodeType>::Insert(const T& element)
{
	// 执行父类插入函数
	BinaryTreeBase<T, NodeType>::Insert(element);

	// 维持树平衡
	OnInsert_MaintainBalance(element);
}

template <typename T, typename NodeType>
void AVLT<T, NodeType>::Delete(Node* deletedNode)
{
	// 先保留被删除节点的父节点
	// 节点被删除之后，对于父节点来说，左右子树的高度会发生变化
	Node* parent = deletedNode->parent;
	
	// 删除节点
	BinaryTreeBase<T, NodeType>::Delete(deletedNode);

	// 如果父节点为空，说明被删除节点是根节点
	if(!parent)
		parent = this->root;
	
	// 维持树平衡
	OnDeleted_MaintainBalance(parent);
}

template <typename T, typename NodeType>
void AVLT<T, NodeType>::OnInsert_MaintainBalance(const T& element)
{
	// 检查二叉树是否平衡
	// 找到最小不平衡子树
	Node* unbalancedSubtree = FindUnbalancedSubtree(this->root, element);

	if(unbalancedSubtree)
	{
		// 不平衡子树根节点的父节点
		Node* unParent = unbalancedSubtree->parent;

		int balanceFactor = this->GetHeightRecursively(unbalancedSubtree->left) - this->GetHeightRecursively(unbalancedSubtree->right);

		// 执行二叉树平衡
		DoBSTBalance(unParent, unbalancedSubtree, element, balanceFactor);
	}
}

template <typename T, typename NodeType>
void AVLT<T, NodeType>::OnDeleted_MaintainBalance(Node* childRoot)
{
	// 检查二叉树是否平衡
	int balanceFactor = this->GetHeightRecursively(childRoot->left) - this->GetHeightRecursively(childRoot->right);
	
	// 如果二叉树不平衡
	if(std::abs(balanceFactor) > 1)
	{
		// 不平衡子树根节点的父节点
		Node* unParent = childRoot->parent;
		Node* deepestLeaf = FindDeepestLeaf(childRoot);

		// 执行二叉树平衡
		DoBSTBalance(unParent, childRoot, deepestLeaf->data, balanceFactor);
	}
}

template <typename T, typename NodeType>
void AVLT<T, NodeType>::DoBSTBalance(Node* parent, Node* unbalancedSubtree, const T& element, const int balanceFactor)
{
	/*
	 * 判断失衡类型
	 * 如果平衡因子大于1，说明左子树比右子树高，那么肯定是LL或者LR的某一种。
	 * 如果平衡因子小于-1，说明右子树比左子树高，那么肯定是RR或者RL的某一种。
	 * 再根据插入值与根节点左孩或者右孩的值的大小关系来判断是L还是R。
	 * 如果更小，那么是L，如果更大，那么是R。
	 */
	Node* newSubtreeRoot = nullptr;
	// LL
	if((balanceFactor > 1) && unbalancedSubtree->left && (element < unbalancedSubtree->left->data))
	{
		// 右旋
		newSubtreeRoot = RightRotate(unbalancedSubtree);
	}
	// LR
	else if((balanceFactor > 1) && unbalancedSubtree->left && (element > unbalancedSubtree->left->data))
	{
		// 先左旋再右旋
		unbalancedSubtree->left = LeftRotate(unbalancedSubtree->left);
		newSubtreeRoot = RightRotate(unbalancedSubtree);
	}
	// RR
	else if((balanceFactor < -1) && unbalancedSubtree->right && (element > unbalancedSubtree->right->data))
	{
		// 左旋
		newSubtreeRoot = LeftRotate(unbalancedSubtree);
	}
	// RL
	else
	{
		// 先右旋再左旋
		unbalancedSubtree->right = RightRotate(unbalancedSubtree->right);
		newSubtreeRoot = LeftRotate(unbalancedSubtree);
	}

	if(parent)
	{
		// 如果子树根的父节点不为空
		// 如果原来的子树是其父节点的左孩
		if(unbalancedSubtree == parent->left)
			parent->left = newSubtreeRoot;
		// 否则为右孩
		else
			parent->right = newSubtreeRoot;
	}
	else
	{
		// 如果子树根的父节点为空，说明当前子树就是整棵树
		// root 指向旋转后的新树的根节点
		this->root = newSubtreeRoot;
	}
}

template <typename T, typename NodeType>
typename AVLT<T, NodeType>::Node* AVLT<T, NodeType>::FindUnbalancedSubtree(Node* childRoot, const T& element)
{
	// 递归树，寻找导致不平衡的插入节点
	if(element < childRoot->data)
	{
		// 递归左子树
		Node* unbalancedSubtree = FindUnbalancedSubtree(childRoot->left, element);
		// 从导致不平衡的子节点返回后，计算当前子树的平衡系数
		int balanceFactor = this->GetHeightRecursively(childRoot->left) - this->GetHeightRecursively(childRoot->right);
		balanceFactor = std::abs(balanceFactor);

		// 如果子树不存在
		if(!unbalancedSubtree)
		{
			// 如果树是平衡的
			if(balanceFactor <= 1)
				unbalancedSubtree = nullptr;
			// 如果树不平衡，说明当前子树就是最小的不平衡子树
			else
				unbalancedSubtree = childRoot;
		}

		return unbalancedSubtree;
	}
	else if(element > childRoot->data)
	{
		// 递归右子树
		Node* unbalancedSubtree = FindUnbalancedSubtree(childRoot->right, element);
		// 从导致不平衡的子节点返回后，计算当前子树的平衡系数
		int balanceFactor = this->GetHeightRecursively(childRoot->left) - this->GetHeightRecursively(childRoot->right);
		balanceFactor = std::abs(balanceFactor);

		// 如果子树不存在
		if(!unbalancedSubtree)
		{
			// 如果树是平衡的
			if(balanceFactor <= 1)
				unbalancedSubtree = nullptr;
			// 如果树不平衡，说明当前子树就是最小的不平衡子树
			else
				unbalancedSubtree = childRoot;
		}

		return unbalancedSubtree;
	}
	else
	{
		// 找到之后，返回到该插入节点的父节点
		return nullptr;
	}
}

template <typename T, typename NodeType>
typename AVLT<T, NodeType>::Node* AVLT<T, NodeType>::LeftRotate(Node* childRoot)
{
	Node* rightChild = childRoot->right;
	Node* temp = rightChild->left;
	
	rightChild->parent = childRoot->parent;
	rightChild->left = childRoot;
	childRoot->parent = rightChild;
	
	childRoot->right = temp;
	if(temp)
		temp->parent = childRoot;

	return rightChild;
}

template <typename T, typename NodeType>
typename AVLT<T, NodeType>::Node* AVLT<T, NodeType>::RightRotate(Node* childRoot)
{
	Node* leftChild = childRoot->left;
	Node* temp = leftChild->right;

	leftChild->parent = childRoot->parent;
	leftChild->right = childRoot;
	childRoot->parent = leftChild;
	
	childRoot->left = temp;
	if(temp)
		temp->parent = childRoot;

	return leftChild;
}

template <typename T, typename NodeType>
typename AVLT<T, NodeType>::Node* AVLT<T, NodeType>::FindDeepestLeaf(Node* childRoot)
{
	int rootHeight = this->GetHeightRecursively(childRoot);

	// 如果不在左子树，那就去找右子树
	Node* result = FindDeepestLeafRecursively(childRoot, childRoot->left, rootHeight);
	if(!result)
		result = FindDeepestLeafRecursively(childRoot, childRoot->right, rootHeight);

	return result;
}

template <typename T, typename NodeType>
typename AVLT<T, NodeType>::Node* AVLT<T, NodeType>::FindDeepestLeafRecursively(Node* origRoot, Node* childRoot, const int rootHeight)
{
	if(!childRoot)
		return nullptr;
	
	Node* result = nullptr;
	// 如果当前节点没有子节点
	if(!(childRoot->left || childRoot->right))
	{
		int depth = CalculateDepth(origRoot, childRoot);
		if(depth == rootHeight)
			result = childRoot;
	}
	else
	{
		// 如果不在左子树，那就去找右子树
		result = FindDeepestLeafRecursively(origRoot, childRoot->left, rootHeight);
		if(!result)
			result = FindDeepestLeafRecursively(origRoot, childRoot->right, rootHeight);
	}

	return result;
}

template <typename T, typename NodeType>
int AVLT<T, NodeType>::CalculateDepth(Node* targetNode, Node* calculatedNode)
{
	if(calculatedNode == targetNode)
	{
		return 0;
	}
	else
	{
		int depth = CalculateDepth(targetNode, calculatedNode->parent);
		return ++depth;
	}
}
